from typing import Any, Protocol, TypedDict

import torch


class TransformerOutput(Protocol):
    tokens_out_emb: torch.Tensor


class Transformer(Protocol):
    def __call__(self, *,
        tokens_in_emb: torch.Tensor,
        tokens_in_pos_emb: torch.Tensor | None = None,
        tokens_in_mask: torch.Tensor | None = None,
        tokens_out_mask: torch.Tensor | None = None,
        **kwargs,
    ) -> TransformerOutput:
        ...


class CacheTransformerOutput(Protocol):
    tokens_out_emb: torch.Tensor
    cache: Any


class CacheTransformer(Transformer):
    def __call__(self, *,
        tokens_in_emb: torch.Tensor,
        tokens_in_pos_emb: torch.Tensor | None = None,
        tokens_in_mask: torch.Tensor | None = None,
        tokens_out_mask: torch.Tensor | None = None,
        cache: Any | None = None,
        **kwargs,
    ) -> CacheTransformerOutput:
        ...
